// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2014 Benoit Steiner <benoit.steiner.goog@gmail.com>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.

#if defined(EIGEN_USE_GPU) && !defined(EIGEN_CXX11_TENSOR_TENSOR_DEVICE_GPU_H)
#define EIGEN_CXX11_TENSOR_TENSOR_DEVICE_GPU_H

// This header file container defines fo gpu* macros which will resolve to
// their equivalent hip* or cuda* versions depending on the compiler in use
// A separate header (included at the end of this file) will undefine all
#include "TensorGpuHipCudaDefines.h"

namespace Eigen {

static const int kGpuScratchSize = 1024;

// This defines an interface that GPUDevice can take to use
// HIP / CUDA streams underneath.
class StreamInterface
{
public:
    virtual ~StreamInterface() {}

    virtual const gpuStream_t& stream() const = 0;
    virtual const gpuDeviceProp_t& deviceProperties() const = 0;

    // Allocate memory on the actual device where the computation will run
    virtual void* allocate(size_t num_bytes) const = 0;
    virtual void deallocate(void* buffer) const = 0;

    // Return a scratchpad buffer of size 1k
    virtual void* scratchpad() const = 0;

    // Return a semaphore. The semaphore is initially initialized to 0, and
    // each kernel using it is responsible for resetting to 0 upon completion
    // to maintain the invariant that the semaphore is always equal to 0 upon
    // each kernel start.
    virtual unsigned int* semaphore() const = 0;
};

class GpuDeviceProperties
{
public:
    GpuDeviceProperties() : initialized_(false), first_(true), device_properties_(nullptr) {}

    ~GpuDeviceProperties()
    {
        if (device_properties_)
        {
            delete[] device_properties_;
        }
    }

    EIGEN_STRONG_INLINE const gpuDeviceProp_t& get(int device) const { return device_properties_[device]; }

    EIGEN_STRONG_INLINE bool isInitialized() const { return initialized_; }

    void initialize()
    {
        if (!initialized_)
        {
            // Attempts to ensure proper behavior in the case of multiple threads
            // calling this function simultaneously. This would be trivial to
            // implement if we could use std::mutex, but unfortunately mutex don't
            // compile with nvcc, so we resort to atomics and thread fences instead.
            // Note that if the caller uses a compiler that doesn't support c++11 we
            // can't ensure that the initialization is thread safe.
            if (first_.exchange(false))
            {
                // We're the first thread to reach this point.
                int num_devices;
                gpuError_t status = gpuGetDeviceCount(&num_devices);
                if (status != gpuSuccess)
                {
                    std::cerr << "Failed to get the number of GPU devices: " << gpuGetErrorString(status) << std::endl;
                    gpu_assert(status == gpuSuccess);
                }
                device_properties_ = new gpuDeviceProp_t[num_devices];
                for (int i = 0; i < num_devices; ++i)
                {
                    status = gpuGetDeviceProperties(&device_properties_[i], i);
                    if (status != gpuSuccess)
                    {
                        std::cerr << "Failed to initialize GPU device #" << i << ": " << gpuGetErrorString(status) << std::endl;
                        gpu_assert(status == gpuSuccess);
                    }
                }

                std::atomic_thread_fence(std::memory_order_release);
                initialized_ = true;
            }
            else
            {
                // Wait for the other thread to inititialize the properties.
                while (!initialized_)
                {
                    std::atomic_thread_fence(std::memory_order_acquire);
                    std::this_thread::sleep_for(std::chrono::milliseconds(1000));
                }
            }
        }
    }

private:
    volatile bool initialized_;
    std::atomic<bool> first_;
    gpuDeviceProp_t* device_properties_;
};

EIGEN_ALWAYS_INLINE const GpuDeviceProperties& GetGpuDeviceProperties()
{
    static GpuDeviceProperties* deviceProperties = new GpuDeviceProperties();
    if (!deviceProperties->isInitialized())
    {
        deviceProperties->initialize();
    }
    return *deviceProperties;
}

EIGEN_ALWAYS_INLINE const gpuDeviceProp_t& GetGpuDeviceProperties(int device) { return GetGpuDeviceProperties().get(device); }

static const gpuStream_t default_stream = gpuStreamDefault;

class GpuStreamDevice : public StreamInterface
{
public:
    // Use the default stream on the current device
    GpuStreamDevice() : stream_(&default_stream), scratch_(NULL), semaphore_(NULL) { gpuGetDevice(&device_); }
    // Use the default stream on the specified device
    GpuStreamDevice(int device) : stream_(&default_stream), device_(device), scratch_(NULL), semaphore_(NULL) {}
    // Use the specified stream. Note that it's the
    // caller responsibility to ensure that the stream can run on
    // the specified device. If no device is specified the code
    // assumes that the stream is associated to the current gpu device.
    GpuStreamDevice(const gpuStream_t* stream, int device = -1) : stream_(stream), device_(device), scratch_(NULL), semaphore_(NULL)
    {
        if (device < 0)
        {
            gpuGetDevice(&device_);
        }
        else
        {
            int num_devices;
            gpuError_t err = gpuGetDeviceCount(&num_devices);
            EIGEN_UNUSED_VARIABLE(err)
            gpu_assert(err == gpuSuccess);
            gpu_assert(device < num_devices);
            device_ = device;
        }
    }

    virtual ~GpuStreamDevice()
    {
        if (scratch_)
        {
            deallocate(scratch_);
        }
    }

    const gpuStream_t& stream() const { return *stream_; }
    const gpuDeviceProp_t& deviceProperties() const { return GetGpuDeviceProperties(device_); }
    virtual void* allocate(size_t num_bytes) const
    {
        gpuError_t err = gpuSetDevice(device_);
        EIGEN_UNUSED_VARIABLE(err)
        gpu_assert(err == gpuSuccess);
        void* result;
        err = gpuMalloc(&result, num_bytes);
        gpu_assert(err == gpuSuccess);
        gpu_assert(result != NULL);
        return result;
    }
    virtual void deallocate(void* buffer) const
    {
        gpuError_t err = gpuSetDevice(device_);
        EIGEN_UNUSED_VARIABLE(err)
        gpu_assert(err == gpuSuccess);
        gpu_assert(buffer != NULL);
        err = gpuFree(buffer);
        gpu_assert(err == gpuSuccess);
    }

    virtual void* scratchpad() const
    {
        if (scratch_ == NULL)
        {
            scratch_ = allocate(kGpuScratchSize + sizeof(unsigned int));
        }
        return scratch_;
    }

    virtual unsigned int* semaphore() const
    {
        if (semaphore_ == NULL)
        {
            char* scratch = static_cast<char*>(scratchpad()) + kGpuScratchSize;
            semaphore_ = reinterpret_cast<unsigned int*>(scratch);
            gpuError_t err = gpuMemsetAsync(semaphore_, 0, sizeof(unsigned int), *stream_);
            EIGEN_UNUSED_VARIABLE(err)
            gpu_assert(err == gpuSuccess);
        }
        return semaphore_;
    }

private:
    const gpuStream_t* stream_;
    int device_;
    mutable void* scratch_;
    mutable unsigned int* semaphore_;
};

struct GpuDevice
{
    // The StreamInterface is not owned: the caller is
    // responsible for its initialization and eventual destruction.
    explicit GpuDevice(const StreamInterface* stream) : stream_(stream), max_blocks_(INT_MAX) { eigen_assert(stream); }
    explicit GpuDevice(const StreamInterface* stream, int num_blocks) : stream_(stream), max_blocks_(num_blocks) { eigen_assert(stream); }
    // TODO(bsteiner): This is an internal API, we should not expose it.
    EIGEN_STRONG_INLINE const gpuStream_t& stream() const { return stream_->stream(); }

    EIGEN_STRONG_INLINE void* allocate(size_t num_bytes) const { return stream_->allocate(num_bytes); }

    EIGEN_STRONG_INLINE void deallocate(void* buffer) const { stream_->deallocate(buffer); }

    EIGEN_STRONG_INLINE void* allocate_temp(size_t num_bytes) const { return stream_->allocate(num_bytes); }

    EIGEN_STRONG_INLINE void deallocate_temp(void* buffer) const { stream_->deallocate(buffer); }

    template <typename Type> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Type get(Type data) const { return data; }

    EIGEN_STRONG_INLINE void* scratchpad() const { return stream_->scratchpad(); }

    EIGEN_STRONG_INLINE unsigned int* semaphore() const { return stream_->semaphore(); }

    EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void memcpy(void* dst, const void* src, size_t n) const
    {
#ifndef EIGEN_GPU_COMPILE_PHASE
        gpuError_t err = gpuMemcpyAsync(dst, src, n, gpuMemcpyDeviceToDevice, stream_->stream());
        EIGEN_UNUSED_VARIABLE(err)
        gpu_assert(err == gpuSuccess);
#else
        EIGEN_UNUSED_VARIABLE(dst);
        EIGEN_UNUSED_VARIABLE(src);
        EIGEN_UNUSED_VARIABLE(n);
        eigen_assert(false && "The default device should be used instead to generate kernel code");
#endif
    }

    EIGEN_STRONG_INLINE void memcpyHostToDevice(void* dst, const void* src, size_t n) const
    {
        gpuError_t err = gpuMemcpyAsync(dst, src, n, gpuMemcpyHostToDevice, stream_->stream());
        EIGEN_UNUSED_VARIABLE(err)
        gpu_assert(err == gpuSuccess);
    }

    EIGEN_STRONG_INLINE void memcpyDeviceToHost(void* dst, const void* src, size_t n) const
    {
        gpuError_t err = gpuMemcpyAsync(dst, src, n, gpuMemcpyDeviceToHost, stream_->stream());
        EIGEN_UNUSED_VARIABLE(err)
        gpu_assert(err == gpuSuccess);
    }

    EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void memset(void* buffer, int c, size_t n) const
    {
#ifndef EIGEN_GPU_COMPILE_PHASE
        gpuError_t err = gpuMemsetAsync(buffer, c, n, stream_->stream());
        EIGEN_UNUSED_VARIABLE(err)
        gpu_assert(err == gpuSuccess);
#else
        eigen_assert(false && "The default device should be used instead to generate kernel code");
#endif
    }

    EIGEN_STRONG_INLINE size_t numThreads() const
    {
        // FIXME
        return 32;
    }

    EIGEN_STRONG_INLINE size_t firstLevelCacheSize() const
    {
        // FIXME
        return 48 * 1024;
    }

    EIGEN_STRONG_INLINE size_t lastLevelCacheSize() const
    {
        // We won't try to take advantage of the l2 cache for the time being, and
        // there is no l3 cache on hip/cuda devices.
        return firstLevelCacheSize();
    }

    EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void synchronize() const
    {
#ifndef EIGEN_GPU_COMPILE_PHASE
        gpuError_t err = gpuStreamSynchronize(stream_->stream());
        if (err != gpuSuccess)
        {
            std::cerr << "Error detected in GPU stream: " << gpuGetErrorString(err) << std::endl;
            gpu_assert(err == gpuSuccess);
        }
#else
        gpu_assert(false && "The default device should be used instead to generate kernel code");
#endif
    }

    EIGEN_STRONG_INLINE int getNumGpuMultiProcessors() const { return stream_->deviceProperties().multiProcessorCount; }
    EIGEN_STRONG_INLINE int maxGpuThreadsPerBlock() const { return stream_->deviceProperties().maxThreadsPerBlock; }
    EIGEN_STRONG_INLINE int maxGpuThreadsPerMultiProcessor() const { return stream_->deviceProperties().maxThreadsPerMultiProcessor; }
    EIGEN_STRONG_INLINE int sharedMemPerBlock() const { return stream_->deviceProperties().sharedMemPerBlock; }
    EIGEN_STRONG_INLINE int majorDeviceVersion() const { return stream_->deviceProperties().major; }
    EIGEN_STRONG_INLINE int minorDeviceVersion() const { return stream_->deviceProperties().minor; }

    EIGEN_STRONG_INLINE int maxBlocks() const { return max_blocks_; }

    // This function checks if the GPU runtime recorded an error for the
    // underlying stream device.
    inline bool ok() const
    {
#ifdef EIGEN_GPUCC
        gpuError_t error = gpuStreamQuery(stream_->stream());
        return (error == gpuSuccess) || (error == gpuErrorNotReady);
#else
        return false;
#endif
    }

private:
    const StreamInterface* stream_;
    int max_blocks_;
};

#if defined(EIGEN_HIPCC)

#define LAUNCH_GPU_KERNEL(kernel, gridsize, blocksize, sharedmem, device, ...)                                \
    hipLaunchKernelGGL(kernel, dim3(gridsize), dim3(blocksize), (sharedmem), (device).stream(), __VA_ARGS__); \
    gpu_assert(hipGetLastError() == hipSuccess);

#else

#define LAUNCH_GPU_KERNEL(kernel, gridsize, blocksize, sharedmem, device, ...)          \
    (kernel)<<<(gridsize), (blocksize), (sharedmem), (device).stream()>>>(__VA_ARGS__); \
    gpu_assert(cudaGetLastError() == cudaSuccess);

#endif

// FIXME: Should be device and kernel specific.
#ifdef EIGEN_GPUCC
static EIGEN_DEVICE_FUNC inline void setGpuSharedMemConfig(gpuSharedMemConfig config)
{
#ifndef EIGEN_GPU_COMPILE_PHASE
    gpuError_t status = gpuDeviceSetSharedMemConfig(config);
    EIGEN_UNUSED_VARIABLE(status)
    gpu_assert(status == gpuSuccess);
#else
    EIGEN_UNUSED_VARIABLE(config)
#endif
}
#endif

}  // end namespace Eigen

// undefine all the gpu* macros we defined at the beginning of the file
#include "TensorGpuHipCudaUndefines.h"

#endif  // EIGEN_CXX11_TENSOR_TENSOR_DEVICE_GPU_H
